/*
 * Copyright (c) 2024 OpenHiTLS. All rights reserved.
 * Licensed under the OpenHiTLS license.
 */
package org.openhitls.tls;

import org.openhitls.core.CryptoNative;
import org.openhitls.core.NativeLoader;

import javax.net.ssl.*;
import java.security.*;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.List;

/**
 * JSSE SSLContext implementation backed by openHiTLS.
 *
 * <p>This class provides TLS context management through the standard
 * Java SSLContext API while using the native openHiTLS library for
 * the actual TLS protocol implementation.</p>
 *
 * <p>Supported protocols:</p>
 * <ul>
 *   <li>TLSv1.2</li>
 *   <li>TLSv1.3</li>
 * </ul>
 */
public class HiTlsSSLContextSpi extends SSLContextSpi {

    // Native SSL_CTX pointer
    private long nativeCtxPtr;

    // Initialization state
    private boolean initialized = false;

    // Key and trust managers
    private X509KeyManager keyManager;
    private X509TrustManager trustManager;
    private SecureRandom secureRandom;

    // Protocol version configuration
    private final int minVersion;
    private final int maxVersion;

    // Supported cipher suites
    private String[] enabledCipherSuites;
    private String[] enabledProtocols;

    // Static initialization
    static {
        NativeLoader.load();
    }

    /**
     * Creates an SSLContext for TLS (default: TLS 1.2 - TLS 1.3).
     */
    public HiTlsSSLContextSpi() {
        this(CryptoNative.TLS_VERSION_1_2, CryptoNative.TLS_VERSION_1_3);
    }

    /**
     * Creates an SSLContext with specified protocol version range.
     *
     * @param minVersion Minimum TLS version
     * @param maxVersion Maximum TLS version
     */
    protected HiTlsSSLContextSpi(int minVersion, int maxVersion) {
        this.minVersion = minVersion;
        this.maxVersion = maxVersion;
    }

    @Override
    protected void engineInit(KeyManager[] km, TrustManager[] tm, SecureRandom sr)
            throws KeyManagementException {
        if (initialized) {
            // Free existing context if re-initializing
            if (nativeCtxPtr != 0) {
                CryptoNative.sslCtxFree(nativeCtxPtr);
                nativeCtxPtr = 0;
            }
        }

        // Extract X509KeyManager
        if (km != null) {
            for (KeyManager manager : km) {
                if (manager instanceof X509KeyManager) {
                    keyManager = (X509KeyManager) manager;
                    break;
                }
            }
        }

        // Extract X509TrustManager
        if (tm != null) {
            for (TrustManager manager : tm) {
                if (manager instanceof X509TrustManager) {
                    trustManager = (X509TrustManager) manager;
                    break;
                }
            }
        }

        // Use provided or default SecureRandom
        this.secureRandom = (sr != null) ? sr : new SecureRandom();

        // Create native SSL context
        try {
            nativeCtxPtr = CryptoNative.sslCtxNew(minVersion, maxVersion);
            if (nativeCtxPtr == 0) {
                throw new KeyManagementException("Failed to create native SSL context");
            }

            // Configure certificates and keys
            configureKeyMaterial();
            configureTrustMaterial();

            // Set default cipher suites and protocols
            initializeDefaults();

            initialized = true;
        } catch (Exception e) {
            if (nativeCtxPtr != 0) {
                CryptoNative.sslCtxFree(nativeCtxPtr);
                nativeCtxPtr = 0;
            }
            throw new KeyManagementException("Failed to initialize SSL context", e);
        }
    }

    private void configureKeyMaterial() throws KeyManagementException {
        if (keyManager == null) {
            return;
        }

        // Get server key alias (or client, depending on usage)
        String alias = keyManager.chooseServerAlias("RSA", null, null);
        if (alias == null) {
            alias = keyManager.chooseServerAlias("EC", null, null);
        }
        if (alias == null) {
            return; // No key available
        }

        // Get certificate chain
        X509Certificate[] chain = keyManager.getCertificateChain(alias);
        if (chain != null && chain.length > 0) {
            try {
                // Set end-entity certificate
                byte[] certDer = chain[0].getEncoded();
                CryptoNative.sslCtxSetCertificate(nativeCtxPtr, certDer, 0); // DER format
            } catch (CertificateEncodingException e) {
                throw new KeyManagementException("Failed to encode certificate", e);
            }
        }

        // Get private key
        PrivateKey privateKey = keyManager.getPrivateKey(alias);
        if (privateKey != null) {
            byte[] keyDer = privateKey.getEncoded();
            if (keyDer != null) {
                CryptoNative.sslCtxSetPrivateKey(nativeCtxPtr, keyDer, 0); // DER format
            }
        }
    }

    private void configureTrustMaterial() throws KeyManagementException {
        if (trustManager == null) {
            // No trust manager provided - use system CA certificates for verification
            // This is the secure default behavior
            CryptoNative.sslCtxEnableVerification(nativeCtxPtr);
            return;
        }

        // Get accepted issuers (CA certificates)
        X509Certificate[] acceptedIssuers = trustManager.getAcceptedIssuers();

        // Check if this is a trust-all manager (empty or null accepted issuers)
        // This is typically used for testing only
        if (acceptedIssuers == null || acceptedIssuers.length == 0) {
            // Trust-all manager detected - enable trust-all mode in native code
            // WARNING: This skips certificate verification - for testing only!
            CryptoNative.sslCtxSetTrustAll(nativeCtxPtr, true);
            return;
        }

        // Normal trust manager with CA certificates
        // Use system CA certificates for verification (custom CA loading is a stub)
        CryptoNative.sslCtxEnableVerification(nativeCtxPtr);

        // Also try to load custom CA certificates from the trust manager
        for (X509Certificate caCert : acceptedIssuers) {
            try {
                byte[] certDer = caCert.getEncoded();
                CryptoNative.sslCtxSetTrustCertificate(nativeCtxPtr, certDer, 0);
            } catch (CertificateEncodingException e) {
                // Log and continue with other certificates
            }
        }

        // Enable peer verification
        CryptoNative.sslCtxSetVerifyMode(nativeCtxPtr, 1);
    }

    private void initializeDefaults() {
        // Initialize enabled cipher suites
        enabledCipherSuites = getDefaultCipherSuites();

        // Initialize enabled protocols based on version range
        List<String> protocols = new ArrayList<>();
        if (minVersion <= CryptoNative.TLS_VERSION_1_2 && maxVersion >= CryptoNative.TLS_VERSION_1_2) {
            protocols.add("TLSv1.2");
        }
        if (maxVersion >= CryptoNative.TLS_VERSION_1_3) {
            protocols.add("TLSv1.3");
        }
        enabledProtocols = protocols.toArray(new String[0]);
    }

    @Override
    protected SSLSocketFactory engineGetSocketFactory() {
        checkInitialized();
        return new HiTlsSSLSocketFactory(this);
    }

    @Override
    protected SSLServerSocketFactory engineGetServerSocketFactory() {
        checkInitialized();
        return new HiTlsSSLServerSocketFactory(this);
    }

    @Override
    protected SSLEngine engineCreateSSLEngine() {
        checkInitialized();
        return new HiTlsSSLEngine(this, null, -1);
    }

    @Override
    protected SSLEngine engineCreateSSLEngine(String host, int port) {
        checkInitialized();
        return new HiTlsSSLEngine(this, host, port);
    }

    @Override
    protected SSLSessionContext engineGetServerSessionContext() {
        checkInitialized();
        // Return a session context for server-side session management
        return new HiTlsSSLSessionContext(true);
    }

    @Override
    protected SSLSessionContext engineGetClientSessionContext() {
        checkInitialized();
        // Return a session context for client-side session management
        return new HiTlsSSLSessionContext(false);
    }

    /**
     * Get the native SSL_CTX pointer.
     * Package-private for use by SSLEngine and Socket implementations.
     */
    long getNativeContext() {
        return nativeCtxPtr;
    }

    /**
     * Get the X509TrustManager for peer verification callbacks.
     */
    X509TrustManager getTrustManager() {
        return trustManager;
    }

    /**
     * Get the X509KeyManager for key selection callbacks.
     */
    X509KeyManager getKeyManager() {
        return keyManager;
    }

    /**
     * Get the enabled cipher suites.
     */
    String[] getEnabledCipherSuites() {
        return enabledCipherSuites.clone();
    }

    /**
     * Set the enabled cipher suites.
     */
    void setEnabledCipherSuites(String[] suites) {
        this.enabledCipherSuites = suites.clone();
        if (nativeCtxPtr != 0) {
            String cipherString = String.join(":", suites);
            CryptoNative.sslCtxSetCipherSuites(nativeCtxPtr, cipherString);
        }
    }

    /**
     * Get the enabled protocols.
     */
    String[] getEnabledProtocols() {
        return enabledProtocols.clone();
    }

    /**
     * Set the enabled protocols.
     */
    void setEnabledProtocols(String[] protocols) {
        this.enabledProtocols = protocols.clone();
    }

    /**
     * Get the supported cipher suites.
     */
    String[] getSupportedCipherSuites() {
        return new String[] {
            // TLS 1.3 cipher suites
            "TLS_AES_128_GCM_SHA256",
            "TLS_AES_256_GCM_SHA384",
            "TLS_CHACHA20_POLY1305_SHA256",
            "TLS_SM4_GCM_SM3",
            "TLS_SM4_CCM_SM3",

            // TLS 1.2 cipher suites
            "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256",
            "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384",
            "TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256",
            "TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384",
            "TLS_RSA_WITH_AES_128_GCM_SHA256",
            "TLS_RSA_WITH_AES_256_GCM_SHA384",

            // SM cipher suites
            "TLS_ECDHE_SM4_CBC_SM3",
            "TLS_ECC_SM4_CBC_SM3"
        };
    }

    /**
     * Get the supported protocols.
     */
    String[] getSupportedProtocols() {
        return new String[] { "TLSv1.2", "TLSv1.3" };
    }

    private String[] getDefaultCipherSuites() {
        return new String[] {
            "TLS_AES_128_GCM_SHA256",
            "TLS_AES_256_GCM_SHA384",
            "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256",
            "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384",
            "TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256",
            "TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384"
        };
    }

    private void checkInitialized() {
        if (!initialized) {
            throw new IllegalStateException("SSLContext has not been initialized");
        }
    }

    @Override
    protected void finalize() throws Throwable {
        try {
            if (nativeCtxPtr != 0) {
                CryptoNative.sslCtxFree(nativeCtxPtr);
                nativeCtxPtr = 0;
            }
        } finally {
            super.finalize();
        }
    }

    // ==================== Inner Classes for TLS version-specific contexts ====================

    /**
     * SSLContext for TLSv1.2 only.
     */
    public static class TLS12Context extends HiTlsSSLContextSpi {
        public TLS12Context() {
            super(CryptoNative.TLS_VERSION_1_2, CryptoNative.TLS_VERSION_1_2);
        }
    }

    /**
     * SSLContext for TLSv1.3 only.
     */
    public static class TLS13Context extends HiTlsSSLContextSpi {
        public TLS13Context() {
            super(CryptoNative.TLS_VERSION_1_3, CryptoNative.TLS_VERSION_1_3);
        }
    }

    /**
     * SSLContext for TLSv1.2 through TLSv1.3.
     */
    public static class TLSContext extends HiTlsSSLContextSpi {
        public TLSContext() {
            super(CryptoNative.TLS_VERSION_1_2, CryptoNative.TLS_VERSION_1_3);
        }
    }
}
